﻿/* Copyright 2008 dnAnalytics Project.
 *
 * Contributors to this file:
 * Marcus Cuda
 *
 * This file is part of dnAnalytics.  dnAnalytics is licensed under the 
 * Microsoft Public License. See License.txt for a complete copy of the
 * license.
 */
using System;
using System.Collections.Generic;
using System.IO;
using System.Text;
using dnAnalytics.LinearAlgebra.IO.Matlab;
using dnAnalytics.Properties;
using zlib;

namespace dnAnalytics.LinearAlgebra.IO
{
    /// <summary>
    /// Writes matrices to a Matlab file.
    /// </summary>
    public class MatlabMatrixWriter : IMatrixWriter, IDisposable
    {
        private const string _headerText = "MATLAB 5.0 MAT-file, Platform: .NET - dnAnalytics, Created on: ";
        private const int _headerTextLength = 116;
        private bool _headerWritten;
        private BinaryWriter _writer;

        /// <summary>
        /// Initializes a new instance of the <see cref="MatlabMatrixWriter"/> class.
        /// </summary>
        /// <param name="filename">The name of the Matlab file to save the matrices to.</param>
        public MatlabMatrixWriter(string filename)
        {
            if (string.IsNullOrEmpty(filename))
            {
                throw new ArgumentException(Resources.EmptyOrNull, "filename");
            }
            _writer = new BinaryWriter(new BufferedStream(new FileStream(filename, FileMode.Create, FileAccess.Write, FileShare.None)));
        }

        #region IDisposable Members

        /// <summary>
        /// Performs application-defined tasks associated with freeing, releasing, or resetting unmanaged resources.
        /// </summary>
        public void Dispose()
        {
            if (_writer != null)
            {
                _writer.Flush();
                _writer.Close();
                _writer = null;
            }
        }

        #endregion

        #region IMatrixWriter Members

        /// <summary>
        /// Writes the given <see cref="Matrix"/> to the file. 
        /// </summary>
        /// <param name="matrix">The matrix to write.</param>
        /// <param name="name">The name of the matrix to store in the file.</param>
        /// <exception cref="ArgumentNullException">If either <paramref name="matrix"/> or <paramref name="name"/> is null.</exception>
        public void WriteMatrix(Matrix matrix, string name)
        {
            if (matrix == null)
            {
                throw new ArgumentNullException("matrix");
            }
            if (string.IsNullOrEmpty(name))
            {
                throw new ArgumentException(Resources.EmptyOrNull, "name");
            }
            if (name.IndexOf(' ') > -1)
            {
                throw new ArgumentException(string.Format(Resources.NameCannotContainASpace, name), "name");
            }

            if (!_headerWritten)
            {
                WriteHeader();
                _headerWritten = true;
            }
            //write datatype
            _writer.Write((int) Matlab.DataType.Compressed);

            if (matrix is SparseMatrix)
            {
                WriteSparseMatrix((SparseMatrix) matrix, name);
            }
            else
            {
                WriteDenseMatrix(matrix, name);
            }
        }

        /// <summary>
        /// Writes the given <see cref="Matrix"/> to the file.
        /// </summary>
        /// <param name="matrices">The matrices to write.</param>
        /// <param name="names">The names of the matrices to store in the file.</param>
        /// <exception cref="ArgumentNullException">If either <paramref name="matrices"/> or <paramref name="names"/> is null.</exception>
        public void WriteMatrices(IList<Matrix> matrices, IList<string> names)
        {
            if (matrices == null)
            {
                throw new ArgumentNullException("matrices");
            }
            if (names == null)
            {
                throw new ArgumentNullException("names");
            }
            if (matrices.Count != names.Count)
            {
                throw new ArgumentException(Resources.ParametersNotConformable);
            }

            for (int i = 0; i < matrices.Count; i++)
            {
                WriteMatrix(matrices[i], names[i]);
            }
        }

        #endregion

        /// <summary>
        /// Closes the stream the being written to.
        /// </summary>
        /// <remarks>Calls <see cref="IDisposable.Dispose"/>.</remarks>
        public void Close()
        {
            Dispose();
        }

        private static void WriteMatrixTagAndName(BinaryWriter writer, ArrayClass arrayClass, string name, int rows, int columns, int nzmax)
        {
            writer.Write((int) Matlab.DataType.Matrix);
            //add place holder for data size
            writer.Write(0);

            //write flag, data type and size
            writer.Write((int) Matlab.DataType.UInt32);
            writer.Write(8);

            //write array class and flags 
            writer.Write((byte) arrayClass);
            writer.Write((byte) 0);

            writer.Write((short) 0);
            writer.Write(nzmax);

            //write dimensions
            writer.Write((int) Matlab.DataType.Int32);
            writer.Write(8);
            writer.Write(rows);
            writer.Write(columns);

            byte[] nameBytes = Encoding.ASCII.GetBytes(name);
            //write name
            if (nameBytes.Length > 4)
            {
                writer.Write((int) Matlab.DataType.Int8);
                writer.Write(nameBytes.Length);
                writer.Write(nameBytes);
                int pad = 8 - (nameBytes.Length%8);
                PadData(writer, pad);
            }
            else
            {
                writer.Write((short) Matlab.DataType.Int8);
                writer.Write((short) nameBytes.Length);
                writer.Write(nameBytes);
                PadData(writer, 4 - nameBytes.Length);
            }
        }

        private void WriteDenseMatrix(Matrix matrix, string name)
        {
            byte[] data = GetDenseDataArray(matrix, name);
            WriteCompressedData(data);
        }

        private void WriteSparseMatrix(SparseMatrix matrix, string name)
        {
            byte[] data = GetSparseDataArray(matrix, name);
            WriteCompressedData(data);
        }

        private static byte[] CommpressData(byte[] data)
        {
            using (MemoryStream compressedStream = new MemoryStream())
            {
                using (ZOutputStream zOutputStream = new ZOutputStream(compressedStream, zlibConst.Z_DEFAULT_COMPRESSION))
                {
                    zOutputStream.Write(data, 0, data.Length);
                }
                return compressedStream.ToArray();
            }
        }

        private static byte[] GetDenseDataArray(Matrix matrix, string name)
        {
            byte[] data;
            using (MemoryStream dataMemoryStream = new MemoryStream())
            using (BinaryWriter dataWriter = new BinaryWriter(dataMemoryStream))
            {
                WriteMatrixTagAndName(dataWriter, ArrayClass.Double, name, matrix.Rows, matrix.Columns, 0);

                //write data
                dataWriter.Write((int) Matlab.DataType.Double);
                dataWriter.Write(matrix.Rows*matrix.Columns*8);

                for (int j = 0; j < matrix.Columns; j++)
                {
                    Vector column = matrix.GetColumn(j);
                    for (int i = 0; i < column.Count; i++)
                    {
                        dataWriter.Write(column[i]);
                    }
                }

                data = dataMemoryStream.ToArray();
            }
            return data;
        }

        private static byte[] GetSparseDataArray(SparseMatrix matrix, string name)
        {
            byte[] data;
            using (MemoryStream dataMemoryStream = new MemoryStream())
            using (BinaryWriter dataWriter = new BinaryWriter(dataMemoryStream))
            {
                int nzmax = matrix.mValueCount;
                WriteMatrixTagAndName(dataWriter, ArrayClass.Sparse, name, matrix.Rows, matrix.Columns, nzmax);

                //write ir
                dataWriter.Write((int) Matlab.DataType.Int32);
                dataWriter.Write(nzmax*4);

                foreach (KeyValuePair<int, Vector> column in matrix.GetColumnEnumerator())
                {
                    foreach (KeyValuePair<int, double> row in column.Value.GetIndexedEnumerator())
                    {
                        dataWriter.Write(row.Key);
                    }
                }

                //add pad if needed
                if (nzmax%2 == 1)
                {
                    dataWriter.Write(0);
                }

                //write jc
                dataWriter.Write((int) Matlab.DataType.Int32);
                dataWriter.Write((matrix.Columns + 1)*4);
                dataWriter.Write(0);
                int count = 0;
                foreach (KeyValuePair<int, Vector> column in matrix.GetColumnEnumerator())
                {
                    count += ((SparseVector) column.Value).mValueCount;
                    dataWriter.Write(count);
                }
                //add pad if needed
                if (matrix.Columns%2 == 0)
                {
                    dataWriter.Write(0);
                }

                //write data
                dataWriter.Write((int) Matlab.DataType.Double);
                dataWriter.Write(nzmax*8);

                foreach (KeyValuePair<int, Vector> column in matrix.GetColumnEnumerator())
                {
                    foreach (KeyValuePair<int, double> row in column.Value.GetIndexedEnumerator())
                    {
                        dataWriter.Write(row.Value);
                    }
                }

                data = dataMemoryStream.ToArray();
            }
            return data;
        }

        private void WriteCompressedData(byte[] data)
        {
            //fill in data size
            byte[] size = BitConverter.GetBytes(data.Length);
            data[4] = size[0];
            data[5] = size[1];
            data[6] = size[2];
            data[7] = size[3];

            //compress data
            byte[] compressedData = CommpressData(data);

            //write compressed data to file
            _writer.Write(compressedData.Length);
            _writer.Write(compressedData);
        }

        private void WriteHeader()
        {
            byte[] header = Encoding.ASCII.GetBytes(_headerText + DateTime.Now.ToString("ddd MMM dd HH:mm:ss yyyy"));
            _writer.Write(header);
            PadData(_writer, _headerTextLength - header.Length + 8, 32);
            //write version
            _writer.Write((short) 0x100);
            //write little endian indicator
            _writer.Write((byte) 0x49);
            _writer.Write((byte) 0x4D);
        }

        private static void PadData(BinaryWriter writer, int bytes)
        {
            PadData(writer, bytes, 0);
        }

        private static void PadData(BinaryWriter writer, int bytes, byte pad)
        {
            for (int i = 0; i < bytes; i++)
            {
                writer.Write(pad);
            }
        }
    }
}